package sse

import (
	"context"
	"errors"
	"fmt"
	"log"
	"mcp-server/internal"
	"time"

	"github.com/mark3labs/mcp-go/client"
	"github.com/mark3labs/mcp-go/mcp"
	"github.com/mark3labs/mcp-go/server"
)

const Version = "1.0.0"

func init() {
	internal.GlobalRegistry.RegisterServer("proxy", &SSEMCPClientConfig{})
}

type SSEMCPClientConfig struct {
	url     string            `json:"url"`
	headers map[string]string `json:"headers"`
	timeout int64             `json:"timeout"`
}

func (c *SSEMCPClientConfig) ParseConfig(config map[string]any) error {
	url, ok := config["url"].(string)
	if !ok {
		return errors.New("missing url")
	}
	c.url = url

	headers, ok := config["headers"].(map[string]string)
	if !ok {
		headers = make(map[string]string)
	}
	c.headers = headers

	timeout, ok := config["timeout"].(int64)
	if !ok {
		timeout = 30
	}
	c.timeout = timeout
	log.Printf("SSEMCPClientConfig ParseConfig: %+v", config)
	return nil
}

func (c *SSEMCPClientConfig) NewServer(serverName string) (*server.MCPServer, error) {
	mcpClient, err := newMCPClient(c)
	if err != nil {
		return nil, fmt.Errorf("Failed to create MCP client: %v", err)
	}

	ctx := context.Background()
	// Start the mcpClient
	if err := mcpClient.Start(ctx); err != nil {
		return nil, fmt.Errorf("Failed to start mcpClient: %v", err)
	}

	mcpServer := server.NewMCPServer(
		serverName,
		Version,
		server.WithResourceCapabilities(true, true),
	)
	info := mcp.Implementation{
		Name:    serverName,
		Version: Version,
	}

	addErr := addClient(ctx, info, mcpClient, mcpServer)
	if addErr != nil {
		return nil, addErr
	}

	return mcpServer, nil
}

func newMCPClient(v *SSEMCPClientConfig) (*client.SSEMCPClient, error) {

	var options []client.ClientOption
	// if v.timeout > 0 {
	// 	options = append(options, client.WithSSEReadTimeout(time.Duration(v.timeout)*time.Second))
	// }
	if len(v.headers) > 0 {
		options = append(options, client.WithHeaders(v.headers))
	}
	return client.NewSSEMCPClient(v.url, options...)

}

func addClient(ctx context.Context, clientInfo mcp.Implementation, mcpClient client.MCPClient, mcpServer *server.MCPServer) error {
	initRequest := mcp.InitializeRequest{}
	initRequest.Params.ProtocolVersion = mcp.LATEST_PROTOCOL_VERSION
	initRequest.Params.ClientInfo = clientInfo
	_, err := mcpClient.Initialize(ctx, initRequest)
	if err != nil {
		return err
	}
	log.Printf("Successfully initialized MCP client")

	err = addClientToolsToServer(ctx, mcpClient, mcpServer)
	if err != nil {
		return err
	}
	_ = addClientPromptsToServer(ctx, mcpClient, mcpServer)
	_ = addClientResourcesToServer(ctx, mcpClient, mcpServer)
	_ = addClientResourceTemplatesToServer(ctx, mcpClient, mcpServer)

	go func() {
		ticker := time.NewTicker(30 * time.Second)
		defer ticker.Stop()
		for {
			select {
			case <-ctx.Done():
				log.Printf("Context done, stopping ping")
				break
			case <-ticker.C:
				_ = mcpClient.Ping(ctx)
			}
		}
	}()
	return nil
}

func addClientToolsToServer(ctx context.Context, mcpClient client.MCPClient, mcpServer *server.MCPServer) error {
	toolsRequest := mcp.ListToolsRequest{}
	for {
		tools, err := mcpClient.ListTools(ctx, toolsRequest)
		if err != nil {
			return err
		}
		log.Printf("Successfully listed %d tools", len(tools.Tools))
		for _, tool := range tools.Tools {
			log.Printf("Adding tool %s", tool.Name)
			mcpServer.AddTool(tool, mcpClient.CallTool)
		}
		if tools.NextCursor == "" {
			break
		}
		toolsRequest.PaginatedRequest.Params.Cursor = tools.NextCursor
	}
	return nil
}

func addClientPromptsToServer(ctx context.Context, mcpClient client.MCPClient, mcpServer *server.MCPServer) error {
	promptsRequest := mcp.ListPromptsRequest{}
	for {
		prompts, err := mcpClient.ListPrompts(ctx, promptsRequest)
		if err != nil {
			return err
		}
		log.Printf("Successfully listed %d prompts", len(prompts.Prompts))
		for _, prompt := range prompts.Prompts {
			log.Printf("Adding prompt %s", prompt.Name)
			mcpServer.AddPrompt(prompt, mcpClient.GetPrompt)
		}
		if prompts.NextCursor == "" {
			break
		}
		promptsRequest.PaginatedRequest.Params.Cursor = prompts.NextCursor
	}
	return nil
}

func addClientResourcesToServer(ctx context.Context, mcpClient client.MCPClient, mcpServer *server.MCPServer) error {
	resourcesRequest := mcp.ListResourcesRequest{}
	for {
		resources, err := mcpClient.ListResources(ctx, resourcesRequest)
		if err != nil {
			return err
		}
		log.Printf("Successfully listed %d resources", len(resources.Resources))
		for _, resource := range resources.Resources {
			log.Printf("Adding resource %s", resource.Name)
			mcpServer.AddResource(resource, func(ctx context.Context, request mcp.ReadResourceRequest) ([]mcp.ResourceContents, error) {
				readResource, e := mcpClient.ReadResource(ctx, request)
				if e != nil {
					return nil, e
				}
				return readResource.Contents, nil
			})
		}
		if resources.NextCursor == "" {
			break
		}
		resourcesRequest.PaginatedRequest.Params.Cursor = resources.NextCursor

	}
	return nil
}

func addClientResourceTemplatesToServer(ctx context.Context, mcpClient client.MCPClient, mcpServer *server.MCPServer) error {
	resourceTemplatesRequest := mcp.ListResourceTemplatesRequest{}
	for {
		resourceTemplates, err := mcpClient.ListResourceTemplates(ctx, resourceTemplatesRequest)
		if err != nil {
			return err
		}
		log.Printf("Successfully listed %d resource templates", len(resourceTemplates.ResourceTemplates))
		for _, resourceTemplate := range resourceTemplates.ResourceTemplates {
			log.Printf("Adding resource template %s", resourceTemplate.Name)
			mcpServer.AddResourceTemplate(resourceTemplate, func(ctx context.Context, request mcp.ReadResourceRequest) ([]mcp.ResourceContents, error) {
				readResource, e := mcpClient.ReadResource(ctx, request)
				if e != nil {
					return nil, e
				}
				return readResource.Contents, nil
			})
		}
		if resourceTemplates.NextCursor == "" {
			break
		}
		resourceTemplatesRequest.PaginatedRequest.Params.Cursor = resourceTemplates.NextCursor
	}
	return nil
}
